Thanks for looking into this, but I don't think this is the right fix.

The current distinction between scheduler types is intentional: timestep-space schedulers call t_to_sigma(), and for flow denoisers hat is where the flow shift is applied. Sigma-space schedulers define their own sigma curve directly from sigma_min() / sigma_max() or from a fixed eference schedule, so applying an additional element-wise flow warp after schedule generation changes the meaning of those schedulers.

This is also not equivalent to diffusers' FlowMatch scheduler behavior with Karras sigmas. Diffusers applies dynamic shifting before the ptional Karras conversion, then Karras rebuilds the schedule from the shifted endpoints. This PR instead warps every generated sigma after the cheduler has already constructed its curve, which changes the intermediate points in a different way.

So I don't think we should merge this global post-processing approach. If we want diffusers-like behavior for Karras, it should be implemented in the scheduler-specific path by shifting the endpoint range before Karras reconstruction, not by warping all scheduler outputs afterward.

I get your point and you're probably right.

But on the other hand, this makes the flow-shift parameter silently innefective with a lot of schedulers, which might feel inconsistent and confusing for users who don't care about the implementation details of the schedulers nor strict compatibility with diffusers.

I understand the concern, but I don't think this UX inconsistency is a problem we should solve at this level.

For most users, schedulers are already black-box sampling strategies. --flow-shift is not guaranteed to have the same type or strength of effect across all schedulers.

For sigma-space schedulers, the schedule itself is the definition of the sigma curve. Applying an element-wise shift after the scheduler has generated that curve changes the scheduler's meaning.

So I think we should keep the current behavior: timestep-space schedulers apply flow shift through t_to_sigma(), while sigma-space schedulers use their own sigma curve, with any shift effect limited to the endpoints when applicable.